import gymnasium as gym
import numpy as np
from numpy import random as rd
import torch
from utils.utils import train_test_split_by_ratio

class MyFuturesEnv(gym.Env):
    def __init__(self, config):
        self.price_ary = config["price_array"].astype(np.float32)
        self.tech_ary = config["tech_array"].astype(np.float32)
        self.if_train = config["if_train"]
        self.train_price_ary, self.train_tech_ary, self.test_price_ary, self.test_tech_ary = train_test_split_by_ratio(
            self.price_ary, self.tech_ary)
        self.future_count = self.price_ary.shape[1]
        self.n_tech = self.tech_ary.shape[1] // self.future_count
        assert self.n_tech * self.future_count == self.tech_ary.shape[1]

        self.gamma = config.get("gamma", 0.99)
        self.min_hand_rate = config.get("min_hand_rate", 0.1)
        self.max_hand = config.get("max_hand", 10)
        self.initial_capital = config.get("initial_capital", 1e8)
        self.buy_cost_pct = config.get("buy_cost_pct", 1e-3)
        self.sell_cost_pct = config.get("sell_cost_pct", 1e-3)
        self.reward_scaling = config.get("reward_scaling", 2 ** 3)
        self.initial_future = (
            np.zeros(self.future_count, dtype=np.int32)
            if config.get("initial_future", None) is None
            else config["initial_future"]
        )
        self.window = config.get("window", 32)
        self.max_position_per_contract = config.get(
            "max_position_per_contract", 100)
        self.N = config.get("N", 5)
        self.forced_liquidation_penalty = config.get(
            "forced_liquidation_penalty", 100)
        self.leverage = config.get("leverage", 10)  # 10倍杠杆，保证金比例10%
        self.bonus_factor = config.get("bonus_factor", 0.1)

        # 状态维度、动作维度等环境元数据
        self.env_name = "MyFuturesEnv"
        self.state_dim = int(
            1 + self.future_count + self.future_count *
            self.window * (1 + self.n_tech)
        )
        self.action_dim = self.future_count
        self.max_step = self.price_ary.shape[0] - 1

        # 初始化资产和回报相关变量
        self.day = None
        self.amount = None  # 账户现金
        self.future = None  # 持仓数量
        self.avg_price = None  # 建仓均价，数组，初始全0
        self.margin_account = None  # 保证金账户余额，数组，初始全0
        self.total_asset = None
        self.gamma_reward = None
        self.initial_total_asset = None
        self.prev_total_asset = None
        self.target_return = 10.0
        self.episode_return = 0.0
        self.float_profit = 0.0

        # 空间定义
        self.observation_space = gym.spaces.Box(
            low=-5000,
            high=5000,
            shape=(self.state_dim,),
            dtype=np.float32,
        )
        self.action_space = gym.spaces.Box(
            low=-1,
            high=1,
            shape=(self.action_dim,),
            dtype=np.float32,
        )

    def reset(self, *, seed=None, options=None):
        self.day = self.window + 1
        price = self.price_ary[self.day]

        if self.if_train:
            while True:
                self.future = (
                        self.initial_future
                        + rd.randint(-16, 16, size=self.initial_future.shape)
                ).astype(np.int32)
                self.amount = (
                        self.initial_capital * rd.uniform(0.95, 1.05)
                        - np.sum(self.future * price / self.leverage)
                )
                if self.amount > 0:
                    break
        else:
            self.future = self.initial_future.astype(np.int32)
            self.amount = self.initial_capital

        # 保证金账户和均价初始化
        self.margin_account = np.abs(self.future * price / self.leverage)
        self.avg_price = np.where(self.future != 0, price, 0.0)

        # 计算浮动盈亏，存为环境成员变量
        self.float_profit = 0.0
        for idx in range(self.future_count):
            pos = self.future[idx]
            self.float_profit += (price[idx] - self.avg_price[idx]) * pos

        # 用现金余额 + 浮动盈亏 计算总资产
        self.total_asset = self.amount + \
                           self.float_profit + sum(self.margin_account)
        self.initial_total_asset = self.total_asset
        self.prev_total_asset = self.total_asset
        self.gamma_reward = 0.0
        self.forced_liquidation = 0  # 强平标志
        self.recent_returns = []

        return self.get_state()

    def step(self, actions):
        actions = (actions * self.max_hand).to(torch.int32).cpu().numpy()
        # actions = (actions * self.max_hand).astype(np.int32)

        price = self.price_ary[self.day]

        min_action = int(self.max_hand * self.min_hand_rate)

        # 开单动作处理
        for idx, act in enumerate(actions):
            price_i = price[idx]
            current_pos = self.future[idx]
            margin_i = self.margin_account[idx]
            avg_p = self.avg_price[idx]

            if act == 0:
                continue

            if act < -min_action:  # 做空动作
                short_volume = -act
                # 平多仓
                if current_pos > 0:
                    # 返还保证金和浮动盈亏到账户现金
                    profit = (price_i - avg_p) * current_pos
                    self.amount += margin_i + profit
                    self.margin_account[idx] = 0
                    self.avg_price[idx] = 0
                    self.future[idx] = 0

                # 开空仓（负仓位）
                max_additional = self.max_position_per_contract + \
                                 self.future[idx]  # future[idx] <=0
                short_volume = min(short_volume, max_additional)
                required_margin = price_i * short_volume / self.leverage
                cost = required_margin * (1 + self.buy_cost_pct)  # 开仓手续费算保证金外

                if self.amount >= cost:
                    # 更新持仓均价（价格平均法）
                    new_pos = self.future[idx] - short_volume
                    if self.future[idx] == 0:
                        self.avg_price[idx] = price_i
                    else:
                        # 计算加权均价，负仓位
                        old_volume = abs(self.future[idx])
                        new_volume = abs(new_pos)
                        self.avg_price[idx] = (
                                                      self.avg_price[idx] * old_volume * 1.0
                                                      + price_i * short_volume
                                              ) / new_volume
                    self.future[idx] = new_pos
                    self.margin_account[idx] += required_margin
                    self.amount -= cost

            elif act > min_action:  # 做多动作
                long_volume = act
                # 平空仓
                if current_pos < 0:
                    profit = (avg_p - price_i) * (-current_pos)
                    self.amount += margin_i + profit
                    self.margin_account[idx] = 0
                    self.avg_price[idx] = 0
                    self.future[idx] = 0

                # 开多仓（正仓位）
                max_additional = self.max_position_per_contract - \
                                 self.future[idx]
                long_volume = min(long_volume, max_additional)
                required_margin = price_i * long_volume / self.leverage
                cost = required_margin * (1 + self.buy_cost_pct)  # 开仓手续费

                if self.amount >= cost:
                    new_pos = self.future[idx] + long_volume
                    if self.future[idx] == 0:
                        self.avg_price[idx] = price_i
                    else:
                        old_volume = abs(self.future[idx])
                        new_volume = abs(new_pos)
                        self.avg_price[idx] = (
                                                      self.avg_price[idx] * old_volume * 1.0
                                                      + price_i * long_volume
                                              ) / new_volume
                    self.future[idx] = new_pos
                    self.margin_account[idx] += required_margin
                    self.amount -= cost

        # 结算
        self.day += 1
        pre_price = price
        price = self.price_ary[self.day]

        # 计算浮动盈亏
        self.float_profit = 0
        for idx in range(self.future_count):
            pos = self.future[idx]
            self.float_profit += (price[idx] - self.avg_price[idx]) * pos

        # 判断爆仓：现金 + 浮动收益总额 < 0
        if (self.amount + self.float_profit) < 0:
            self.forced_liquidation = 1
            # 平所有仓位，返还保证金和浮动盈亏到账户现金
            for idx in range(self.future_count):
                pos = self.future[idx]
                if pos != 0:
                    profit = 0
                    if pos > 0:
                        profit = (price[idx] - self.avg_price[idx]) * pos
                    else:
                        profit = (self.avg_price[idx] - price[idx]) * (-pos)
                    self.amount += self.margin_account[idx] + profit
                    self.margin_account[idx] = 0
                    self.avg_price[idx] = 0
                    self.future[idx] = 0
        else:
            self.forced_liquidation = 0

        if self.forced_liquidation:
            self.float_profit = 0.0
        else:
            # 否则按当前仓位重新计算浮动盈亏
            self.float_profit = 0.0
            for idx in range(self.future_count):
                pos = self.future[idx]
                self.float_profit += (price[idx] - self.avg_price[idx]) * pos

        self.total_asset = self.amount + \
                           self.float_profit + sum(self.margin_account)
        state = self.get_state()

        price_change = (price - pre_price) / np.maximum(pre_price, 1e-9)
        reward = self.get_reward(
            total_asset=self.total_asset,
            action=actions,
            price_change=price_change,
        )
        # reward = self.compute_reward(self.total_asset)

        self.gamma_reward = self.gamma_reward * self.gamma + reward
        done = self.day == self.max_step or self.total_asset < self.initial_total_asset * 0.05
        if done:
            self.episode_return = self.total_asset / self.initial_total_asset

        self.prev_total_asset = self.total_asset

        return state, reward, done, dict()

    # 暂时不起用，用get_reward
    def compute_reward(self, total_asset):
        prob_reward = (total_asset - self.prev_total_asset) / \
                      max(self.prev_total_asset, 1e-9)

        self.recent_returns.append(prob_reward)
        if len(self.recent_returns) > self.N:
            self.recent_returns.pop(0)
        downside = [r for r in self.recent_returns if r < 0]
        downside_risk = np.std(downside) if downside else 0.0
        forced_penalty = self.forced_liquidation * self.forced_liquidation_penalty
        reward = prob_reward - 0.5 * downside_risk - forced_penalty
        return reward

    def get_reward(self, total_asset, action, price_change):
        profit_reward = (total_asset - self.prev_total_asset) / max(self.prev_total_asset, 1e-9)

        penalty_liquidation = self.forced_liquidation * self.forced_liquidation_penalty

        # -------- 顺势操作奖励--------
        trend_alignment = action * price_change
        trend_reward = self.bonus_factor * np.mean(trend_alignment)  # 奖励系数
        # -------- 下行风险惩罚 --------
        self.recent_returns.append(profit_reward)
        if len(self.recent_returns) > self.N:
            self.recent_returns.pop(0)
        downside = [r for r in self.recent_returns if r < 0]
        downside_risk = np.std(downside) if downside else 0.0
        risk_penalty = 0.5 * downside_risk

        # --------最终 reward 汇总--------
        reward = (
                profit_reward
                + trend_reward
                - penalty_liquidation
                - risk_penalty
        )
        return reward * self.reward_scaling

    def get_state(self):
        available_amount = np.array(
            self.amount + self.float_profit, dtype=np.float32)
        start = max(0, self.day - self.window)
        end = self.day
        price_window = self.price_ary[start:end].reshape(-1)
        tech_window = self.tech_ary[start:end].reshape(-1)
        return np.hstack(
            (
                available_amount / self.initial_capital,
                self.future.astype(np.float32) /
                self.max_position_per_contract,
                price_window,
                tech_window,
            )
        ).astype(np.float32)

    def get_window(self):
        return int(self.window)

    def get_future_count(self):
        return int(self.future_count)

    def get_tech_num(self):
        return int(self.n_tech)

    def get_day(self):
        return int(self.day)

    def get_max_step(self):
        return int(self.max_step)

    def get_final_ratio(self):
        return self.total_asset / self.initial_total_asset

    @staticmethod
    def sigmoid_sign(ary, thresh):
        def sigmoid(x):
            return 1 / (1 + np.exp(-x * np.e)) - 0.5

        return sigmoid(ary / thresh) * thresh

    def set_train(self):
        self.if_train = True
        self.price_ary, self.tech_ary = self.train_price_ary, self.train_tech_ary
        self.max_step = self.price_ary.shape[0] - 1

    def set_test(self):
        self.if_train = False
        self.price_ary, self.tech_ary = self.test_price_ary, self.test_tech_ary
        self.max_step = self.price_ary.shape[0] - 1

    def get_price_array(self):
        return self.price_ary[self.window :-1]


